/**
 * 
 */
package edu.umd.clip.lm.programs;

import edu.berkeley.nlp.util.*;

import java.io.*;
import java.net.*;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import edu.umd.clip.lm.model.*;
import edu.umd.clip.lm.model.decoding.NgramMarginalization;
import edu.umd.clip.jobs.*;
import edu.umd.clip.lm.factors.*;
import edu.umd.clip.lm.util.*;

import java.lang.management.*;

import com.sleepycat.je.DatabaseException;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class NgramServer {
	public static class Options {
        @Option(name = "-config", required = true, usage = "XML config file")
		public String config;
        @Option(name = "-jobs", usage = "number of concurrent jobs (default: 1)")
        public int jobs = 1;
        @Option(name = "-forest", required = true, usage = "the decision tree forest")
		public String forest;        
        @Option(name = "-debug", usage = "decoder debug level (default: 0)")
        public int debug = 0;
        @Option(name = "-host", usage = "remote storage hostname or IP address (default: use local storage)")
        public String host = null;
        @Option(name = "-port", usage = "remove storage port (default: 2332)")
        public int port = 0;
        @Option(name = "-listen-port", required = true, usage = "port to listen at")
        public int listenPort = 0;
        @Option(name = "-meminfo", usage = "track memory usage (default: false)")
        public boolean meminfo = false;
        @Option(name = "-coarse-threshold", usage = "coarse threshold, requires the new decoder (default: 1e-4)")
        public double coarse = -1;
        @Option(name = "-lc", usage = "lowercase the input")
        public boolean lower = false;
        @Option(name = "-uc", usage = "uppercase the input")
        public boolean upper = false;
        @Option(name = "-jerboa", required = false, usage = "use Jerboa storage (default: false)")
		public boolean useJerboa = false;        
        @Option(name = "-compact", required = false, usage = "use Compact storage (default: false)")
		public boolean useCompact = false;        
        @Option(name = "-bdb", required = false, usage = "use Berkeley DB storage (default: false)")
		public boolean useBDB = false;        
        @Option(name = "-srilm", required = false, usage = "use SRILM protocol (default: false)")
		public boolean srilmProtocol = false;        
	}

	/**
	 * @param args
	 * @throws DatabaseException 
	 * @throws IOException 
	 * @throws InterruptedException 
	 * @throws ClassNotFoundException 
	 */
	public static void main(String[] args) throws DatabaseException, IOException, InterruptedException, ClassNotFoundException {
		final MemoryMXBean memoryBean = ManagementFactory.getMemoryMXBean();
		
		MemoryUsage memuse;
		
        OptionParser optParser = new OptionParser(Options.class);
        final Options opts = (Options) optParser.parse(args, true);

        if (opts.upper && opts.lower) {
        	System.err.println("Lowercase AND uppercase? doesn't sound right.");
        	return;
        }

        if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: initial: %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
		}

        LMDecodingOptions lmOpts = new LMDecodingOptions();
        lmOpts.config = opts.config;
        lmOpts.debug = opts.debug;
        if (opts.jobs > 0) lmOpts.jobs = opts.jobs;
        if (opts.host != null) lmOpts.host = opts.host;
        if (opts.port > 0) lmOpts.port = opts.port;
        if (opts.coarse > 0) lmOpts.coarseThreshold = opts.coarse;
        if (opts.forest != null) lmOpts.forest = opts.forest;
        
        if (opts.host != null && opts.port > 0) lmOpts.storage = LMDecodingOptions.Storage.REMOTE;
        if (opts.useJerboa) lmOpts.storage = LMDecodingOptions.Storage.JERBOA;
        if (opts.useBDB) lmOpts.storage = LMDecodingOptions.Storage.BDB;
        if (opts.useCompact) lmOpts.storage = LMDecodingOptions.Storage.COMPACT;
        
		if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: initial: %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
		}

        LanguageModel.initDecoding(lmOpts);

        Experiment experiment = Experiment.getInstance();

		if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: after Experiment.initialize(): %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
		}
		
		final InputParser parser = new PlainInputParser(experiment.getTupleDescription());
		/*	
		{
			long[] endTuples = new long[1];
			endTuples[0] = experiment.getTupleDescription().createEndTuple(); 
			parser.setEndTuples(endTuples);
		}
		*/
		final ForestModel forest = experiment.getForest(opts.forest);
		
		if (opts.meminfo) {
			memuse = memoryBean.getHeapMemoryUsage();
			System.out.printf("MEMUSE: after loading LMs: %dM/%dM\n", memuse.getUsed()/1048576, memuse.getMax()/1048576);
		}

		final NgramMarginalization marginalizer = new NgramMarginalization(forest);
		
		ServerSocket server = new ServerSocket(opts.listenPort);
		
		while(true) {
			final Socket socket = server.accept();
			socket.setTcpNoDelay(true);
			
			Runnable run = new Runnable() {

				@Override
				public void run() {
					try {
						NgramProtocolHandler handler;
						if (opts.srilmProtocol) {
							handler = new SRILMNgramProtocolHandler(socket, marginalizer, opts);
						} else {
							handler = new SimpleNgramProtocolHandler(socket, marginalizer, opts);
						}

						handler.processInput();
						handler.finish();
						
					    if (opts.debug > 0) {
					    	System.err.println(forest.getName() + " stats: \n" + forest.getDecoder().getStatistics() + "\n");
					    }
					} catch(IOException e) {
						e.printStackTrace();
					} catch (InterruptedException e) {
						e.printStackTrace();
					}
					try {
						socket.close();
					} catch (IOException e) {}
				}
				
			};
			
			new Thread(run, "serving "+socket.getInetAddress().toString()).start();
		}
	}
	
	private abstract static class NgramProtocolHandler {
		final BufferedReader reader;
		final BufferedWriter writer;
		final NgramMarginalization marginalizer;
		final Options opts;
		final InputParser parser;
		
		NgramProtocolHandler(Socket socket, NgramMarginalization marginalizer, Options opts) throws IOException {
			this.marginalizer = marginalizer;
			this.opts = opts;
			reader = IO.getReader(socket.getInputStream());
		    writer = IO.getWriter(socket.getOutputStream());
			parser = new PlainInputParser(Experiment.getInstance().getTupleDescription());
		}
		
		public abstract void processInput() throws IOException;
		
		public void finish() throws IOException, InterruptedException {
			reader.close();
			writer.close();
		}
	}
	
	private static class SRILMNgramProtocolHandler extends NgramProtocolHandler {
		private static final String REMOTELM_VERSION2 = "_R_E_M_O_T_E_L_M_V=2";
		private static final String REMOTELM_WORDPROB = "W";
		private static final String REMOTELM_CONTEXTID1	= "C1";
		private static final String REMOTELM_CONTEXTID2	= "C2";
		private static final String REMOTELM_CONTEXTBOW	= "B";
		private static final String REMOTELM_OK = "OK";
		private static final String REMOTELM_ERROR = "ERROR";
		private static final Pattern lineRE = Pattern.compile("^(\\w+)\\s+(.*)");
		
		private int protocolVersion;
		/**
		 * @param socket
		 * @param marginalizer
		 * @param opts
		 * @throws IOException
		 */
		SRILMNgramProtocolHandler(Socket socket,
				NgramMarginalization marginalizer, Options opts)
				throws IOException 
		{
			super(socket, marginalizer, opts);
			protocolVersion = 1;
			
			writer.append("probserver ready\n");
			writer.flush();
		}

		/* (non-Javadoc)
		 * @see edu.umd.clip.lm.programs.NgramServer.NgramProtocolHandler#processInput()
		 */
		@Override
		public void processInput() throws IOException {
		    long lineNo = 0;
		    for(String line = reader.readLine(); line != null; line = reader.readLine()) {
		    	++lineNo;
		    	//System.err.printf("SRI input[%06d]: %s\n", lineNo, line);
		    	
		    	line = line.trim();
		    	
		    	if (line.startsWith(REMOTELM_VERSION2)) {
		    		protocolVersion = 2;
		    		writer.write(REMOTELM_OK);
		    		writer.write("\n");
		    	} else if (protocolVersion == 1 || line.startsWith(REMOTELM_WORDPROB)) {
			    	long[] ngram = getNgram(line);
		    		if (ngram != null) {
				    	double logProb = marginalizer.computeNgramProbability(ngram);
				    	if (protocolVersion == 1) {
				    		writer.write(String.format("%f\n", logProb));				    		
				    	} else {
				    		writer.write(String.format("%s %f\n", REMOTELM_OK, logProb));
				    	}
		    		} else {
		    			writer.write(REMOTELM_ERROR);
		    			writer.write("\n");
		    		}
		    	} else if (line.startsWith(REMOTELM_CONTEXTBOW)) {
		    		writer.write(String.format("%s %f\n", REMOTELM_OK, 0.0f));
		    	} else if (line.startsWith(REMOTELM_CONTEXTID1)) {
		    		long[] ngram = getNgram(line);
		    		int hash = Arrays.hashCode(ngram);
		    		writer.write(String.format("%s %d %d\n", REMOTELM_OK, hash, ngram.length));
		    	} else if (line.startsWith(REMOTELM_CONTEXTID2)) {
		    		long[] ngram = getNgram(line);
		    		int hash = Arrays.hashCode(Arrays.copyOf(ngram, ngram.length-1));
		    		writer.write(String.format("%s %d %d\n", REMOTELM_OK, hash, ngram.length-1));		    		
		    	} 
		    	writer.flush();
		    }
		}
		
		private long[] getNgram(String line) {
    		Matcher matcher = lineRE.matcher(line);
    		if (matcher.matches()) {
    			String ngramStr = matcher.group(2);
    			
				if (opts.lower) {
					ngramStr = ngramStr.toLowerCase();
				} else if (opts.upper) {
					ngramStr = ngramStr.toUpperCase();
				}
				if (opts.debug > 0) {
					System.err.printf("ngram input: '%s'\n", ngramStr);
				}
		    	long[] ngram = parser.parseSentence(ngramStr);
		    	return ngram;
    		}
    		return null;
		}
	}
	
	private static class SimpleNgramProtocolHandler extends NgramProtocolHandler {
		final DataSequencer<String> sequencer;
		final Thread outputPrinter;

		/**
		 * @param socket
		 * @param marginalizer
		 * @throws IOException
		 */
		SimpleNgramProtocolHandler(Socket socket, NgramMarginalization marginalizer, Options opts) throws IOException {
			super(socket, marginalizer, opts);
		    sequencer = new DataSequencer<String>(1);
		    outputPrinter = new Thread(new OutputPrinter(sequencer, writer));
		}

		/* (non-Javadoc)
		 * @see edu.umd.clip.lm.programs.NgramServer.NgramProtocolHandler#processInput()
		 */
		@Override
		public void processInput() throws IOException {
		    JobManager manager = JobManager.getInstance();
		    JobGroup group = manager.createJobGroup("decoding");
		    
		    outputPrinter.start();
		    
		    int lineNo = 0;
		    for(String line = reader.readLine(); line != null; line = reader.readLine()) {
		    	Job job = new Job(new EvalRunnable(++lineNo, line), "decoding sentence #"+Integer.toString(lineNo));
		    	manager.addJob(group, job);
		    }
		    sequencer.setLastItem(lineNo);
		    try {
				outputPrinter.join();
			} catch (InterruptedException e) {
				e.printStackTrace();
			}
		}
		
		private class EvalRunnable implements Runnable {
			final Pair<Integer, String> pair;
			
			public EvalRunnable(int sentNo, String line) {
				this.pair = new Pair<Integer, String>(sentNo, line);
			}
			
			public void run() {
				String input = pair.getSecond();
				if (opts.lower) {
					input = input.toLowerCase();
				} else if (opts.upper) {
					input = input.toUpperCase();
				}
				if (opts.debug > 0) {
					System.err.printf("ngram input: '%s'\n", input);
				}
		    	long[] ngram = parser.parseSentence(input);
		    	
		    	double logProb = marginalizer.computeNgramProbability(ngram);
		    	
				StringBuilder sb = new StringBuilder();
		    	sb.append(pair.getSecond());
		    	sb.append("\t");
		    	sb.append(Float.toString((float)logProb));

		    	pair.setSecond(sb.toString());
		    	
		    	//System.err.printf("done line %d, %s\n", pair.getFirst(), pair.getSecond());
		    	//System.err.flush();

		    	sequencer.putItem(pair);
			}
		}

	}
	
	private static class OutputPrinter implements Runnable {
		private final DataSequencer<String> sequencer;
		private final BufferedWriter writer;
		/**
		 * @param sequencer
		 * @param writer
		 */
		public OutputPrinter(DataSequencer<String> sequencer,
				BufferedWriter writer) {
			this.sequencer = sequencer;
			this.writer = writer;
		}

		/* (non-Javadoc)
		 * @see java.lang.Runnable#run()
		 */
		@Override
		public void run() {
		    try {
			    while(!sequencer.isDone()) {
			    	Pair<Integer, String> item = sequencer.getFirst();
			    	if (item== null) return;
			    	//System.err.printf("printing line %d, %s\n", item.getFirst(), item.getSecond());
			    	//System.err.flush();
			    	writer.append(item.getSecond());
			    	writer.newLine();
			    	writer.flush();
			    }
		    } catch(InterruptedException e) {
		    	e.printStackTrace();
		    } catch (IOException e) {
				e.printStackTrace();
			}
			
		}
		
	}
}
